import whisper
import time
import tkinter as tk
from tkinter import ttk, filedialog, messagebox, scrolledtext
from threading import Thread
import os
import sys
import torch
from pathlib import Path

# 支持的音视频格式
SUPPORTED_EXTS = {'.wav', '.mp3', '.m4a', '.ogg', '.mp4', '.mov', '.avi', '.flv'}


class BatchAudioProcessor:
    def __init__(self, root):
        self.root = root
        self.root.title("批量音视频处理工具（显存优化版）")
        self.root.geometry("1200x900")

        # 核心状态变量
        self.whisper_model = None  # Whisper模型句柄（每次处理后释放）
        self.ollama_model = None  # Ollama模型句柄（每次处理后释放）
        self.target_dir = ""  # 选中的文件夹路径
        self.file_queue = []  # 待处理文件队列
        self.current_file_idx = 0  # 当前处理文件索引
        self.processing = False  # 批量处理状态标记

        # 自动保存配置
        self.output_dir = "batch_outputs"
        os.makedirs(self.output_dir, exist_ok=True)

        # 创建界面
        self.create_widgets()

    def create_widgets(self):
        """创建界面组件（新增显存监控标签）"""
        # 顶部标题
        ttk.Label(
            self.root,
            text="批量音视频处理工具（显存优化版）",
            font=("微软雅黑", 24, "bold"),
            padding=20
        ).pack(fill=tk.X)

        # 显存状态监控（新增）
        self.lbl_gpu_memory = ttk.Label(
            self.root,
            text="显存状态: 未初始化",
            foreground="blue",
            padding=5
        )
        self.lbl_gpu_memory.pack(fill=tk.X)

        # 文件夹选择区域
        frame_dir = ttk.LabelFrame(self.root, text="选择处理文件夹", padding=15)
        frame_dir.pack(fill=tk.X, pady=10, padx=20)

        ttk.Button(
            frame_dir,
            text="选择音视频文件夹",
            command=self.select_folder,
            width=25
        ).grid(row=0, column=0, padx=5, pady=5)

        self.lbl_dir_path = ttk.Label(
            frame_dir,
            text="未选择文件夹",
            wraplength=800,
            anchor=tk.W
        )
        self.lbl_dir_path.grid(row=0, column=1, padx=5, pady=5, sticky=tk.W)

        # 文件列表显示区域
        frame_files = ttk.LabelFrame(self.root, text="待处理文件列表", padding=15)
        frame_files.pack(fill=tk.BOTH, expand=True, pady=10, padx=20)

        self.txt_file_list = scrolledtext.ScrolledText(
            frame_files,
            wrap=tk.WORD,
            height=8,
            font=("微软雅黑", 12),
            width=100
        )
        self.txt_file_list.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

        # 处理控制区域
        frame_controls = ttk.Frame(self.root, padding=15)
        frame_controls.pack(fill=tk.X, pady=10)

        self.btn_start = ttk.Button(
            frame_controls,
            text="开始批量处理",
            command=self.start_batch_processing,
            width=25,
            state=tk.DISABLED
        )
        self.btn_start.pack(side=tk.LEFT, padx=10)

        self.lbl_progress = ttk.Label(
            frame_controls,
            text="状态：等待选择文件夹",
            font=("微软雅黑", 12)
        )
        self.lbl_progress.pack(side=tk.LEFT, padx=20)

        # 输出日志区域
        frame_log = ttk.LabelFrame(self.root, text="处理日志", padding=15)
        frame_log.pack(fill=tk.BOTH, expand=True, pady=10, padx=20)

        self.txt_log = scrolledtext.ScrolledText(
            frame_log,
            wrap=tk.WORD,
            height=10,
            font=("微软雅黑", 12),
            width=100
        )
        self.txt_log.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)

    def select_folder(self):
        """选择文件夹并加载待处理文件"""
        folder_path = filedialog.askdirectory(title="选择包含音视频文件的文件夹")
        if not folder_path:
            return

        self.target_dir = folder_path
        self.lbl_dir_path.config(text=f"已选择文件夹: {folder_path}")

        # 扫描文件夹并筛选支持的文件
        self.file_queue = []
        for file in Path(folder_path).rglob('*'):
            if file.is_file() and file.suffix.lower() in SUPPORTED_EXTS:
                self.file_queue.append(str(file))

        # 显示文件列表
        self.txt_file_list.delete(1.0, tk.END)
        if self.file_queue:
            self.txt_file_list.insert(tk.END, "待处理文件（共{}个）:\n".format(len(self.file_queue)))
            for idx, file in enumerate(self.file_queue, 1):
                self.txt_file_list.insert(tk.END, f"{idx}. {file}\n")
            self.btn_start.config(state=tk.NORMAL)
            self.lbl_progress.config(text=f"状态：发现{len(self.file_queue)}个待处理文件")
        else:
            messagebox.showwarning("提示", "该文件夹中未找到支持的音视频文件")
            self.btn_start.config(state=tk.DISABLED)

    def get_gpu_memory(self):
        """获取当前显存占用（仅GPU有效）"""
        if not torch.cuda.is_available():
            return "无GPU可用"
        allocated = torch.cuda.memory_allocated() / 1024 ** 2  # MB
        cached = torch.cuda.memory_reserved() / 1024 ** 2
        return f"已分配: {allocated:.2f}MB | 缓存: {cached:.2f}MB"

    def update_gpu_status(self):
        """更新显存状态标签"""
        self.root.after(0, lambda: self.lbl_gpu_memory.config(
            text=f"显存状态: {self.get_gpu_memory()}"
        ))

    def load_whisper_model(self):
        """加载Whisper模型（临时加载，用完释放）"""
        try:
            self.whisper_model = whisper.load_model("turbo", in_memory=True)
            self.log("Whisper模型加载成功")
            self.update_gpu_status()
            return True
        except Exception as e:
            self.log(f"Whisper模型加载失败: {str(e)}")
            return False

    def unload_whisper_model(self):
        """卸载Whisper模型并释放显存"""
        if self.whisper_model:
            del self.whisper_model
            self.whisper_model = None
            torch.cuda.empty_cache()  # 清理未释放的缓存
            torch.cuda.ipc_collect()  # 收集碎片化内存
            self.log("Whisper模型已卸载，显存已释放")
            self.update_gpu_status()

    def load_ollama_model(self):
        """加载Ollama模型（临时加载，用完释放）"""
        try:
            import ollama
            self.ollama_model = ollama  # 保存模块引用（实际Ollama是客户端，无显式模型对象）
            self.log("Ollama服务连接成功")
            self.update_gpu_status()
            return True
        except Exception as e:
            self.log(f"Ollama连接失败: {str(e)}")
            return False

    def unload_ollama_model(self):
        """卸载Ollama相关资源（释放隐性显存占用）"""
        # Ollama通过HTTP调用，无显式模型对象，但可通过gc清理
        import gc
        gc.collect()  # 强制垃圾回收
        torch.cuda.empty_cache()
        self.log("Ollama资源已释放")
        self.update_gpu_status()

    def start_batch_processing(self):
        """启动批量处理（增加显存初始化检查）"""
        if self.processing:
            messagebox.showwarning("提示", "正在处理中，请耐心等待")
            return

        # 初始化时检查GPU
        if torch.cuda.is_available():
            self.log(f"检测到GPU: {torch.cuda.get_device_name(0)}")
        else:
            self.log("当前使用CPU模式")

        self.processing = True
        self.current_file_idx = 0
        self.btn_start.config(text="处理中...", state=tk.DISABLED)
        self.log("\n===== 开始批量处理 =====")

        Thread(target=self.process_files_in_queue, daemon=True).start()

    def process_files_in_queue(self):
        """处理队列中的每个文件（严格释放每一步显存）"""
        for idx, file_path in enumerate(self.file_queue, 1):
            self.current_file_idx = idx
            self.update_progress(f"处理第 {idx}/{len(self.file_queue)} 个文件: {os.path.basename(file_path)}")

            try:
                # 步骤1：加载Whisper→转录→卸载Whisper
                if not self.load_whisper_model():
                    continue
                trans_text = self.transcribe_audio(file_path)
                self.unload_whisper_model()  # 转录完成立即卸载

                if not trans_text:
                    self.log(f"⚠️ 转录失败: {file_path}")
                    continue

                # 步骤2：加载Ollama→生成摘要→卸载Ollama
                if not self.load_ollama_model():
                    continue
                llm_response = self.generate_summary(trans_text)
                self.unload_ollama_model()  # 生成完成立即卸载

                if not llm_response:
                    self.log(f"⚠️ 对话生成失败: {file_path}")
                    continue

                # 步骤3：保存结果
                save_path = self.save_result(file_path, trans_text, llm_response)
                self.log(f"✅ 处理完成: {save_path}")

            except Exception as e:
                self.log(f"❌ 处理异常: {file_path} - {str(e)}")
                continue

        # 全部处理完成
        self.root.after(0, lambda: self.batch_complete())

    def transcribe_audio(self, file_path):
        """转录单个音频文件（增加显存监控）"""
        self.log(f"开始转录: {file_path}")
        start_time = time.time()
        self.update_gpu_status()  # 记录转录前显存

        try:
            result = self.whisper_model.transcribe(
                file_path,
                #language="zh",
                fp16=torch.cuda.is_available()
            )
            trans_time = time.time() - start_time
            self.log(f"转录完成（耗时{trans_time:.2f}秒），文本长度: {len(result['text'])}字")
            self.update_gpu_status()  # 记录转录后显存
            return result['text']

        except Exception as e:
            self.log(f"转录错误: {str(e)}")
            return None

    def generate_summary(self, text):
        """调用大模型生成摘要（增加显存监控）"""
        self.log("调用大模型生成摘要...")
        self.update_gpu_status()  # 记录生成前显存

        try:
            response = self.ollama_model.generate('qwen3', f"整理这个视频内容，提取出核心内容：{text}")
            self.log("摘要生成成功")
            self.update_gpu_status()  # 记录生成后显存
            return response['response']

        except Exception as e:
            self.log(f"大模型调用失败: {str(e)}")
            return None

    def save_result(self, original_path, trans_text, llm_response):
        """保存结果（逻辑不变）"""
        base_name = os.path.splitext(os.path.basename(original_path))[0]
        save_dir = os.path.join(self.output_dir, base_name)
        os.makedirs(save_dir, exist_ok=True)

        with open(os.path.join(save_dir, f"{base_name}_transcription.txt"), 'w', encoding='utf-8') as f:
            f.write(trans_text)

        with open(os.path.join(save_dir, f"{base_name}_summary.txt"), 'w', encoding='utf-8') as f:
            f.write(llm_response)

        return save_dir

    def update_progress(self, msg):
        """更新进度提示"""
        self.root.after(0, lambda: self.lbl_progress.config(text=f"状态：{msg}"))
        self.log(msg)

    def log(self, msg):
        """记录日志到文本框"""
        self.root.after(0, lambda: self.txt_log.insert(tk.END, f"{time.strftime('%H:%M:%S')} {msg}\n"))
        self.root.after(0, lambda: self.txt_log.see(tk.END))

    def batch_complete(self):
        """批量处理完成后的状态重置"""
        self.processing = False
        self.btn_start.config(text="开始批量处理", state=tk.NORMAL)
        self.log("\n===== 批量处理完成 =====")
        messagebox.showinfo("完成",
                            f"批量处理完成！共处理{len(self.file_queue)}个文件，结果保存在: {os.path.abspath(self.output_dir)}")


if __name__ == "__main__":
    root = tk.Tk()
    app = BatchAudioProcessor(root)
    root.mainloop()
